package com.zyx.flinkdemo.stream.state.keyed;

import org.apache.flink.api.common.accumulators.AverageAccumulator;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.AggregatingState;
import org.apache.flink.api.common.state.AggregatingStateDescriptor;
import org.apache.flink.api.java.tuple.Tuple;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author Yaxi.Zhang
 * @since 2021/6/2 17:39
 * desc: AggregatingState案例
 */

public class AggregatingStateFlatMap extends RichFlatMapFunction<Tuple2<Integer, Integer>, Tuple2<Integer, Double>> {

    public static Logger log = LoggerFactory.getLogger(AggregatingStateFlatMap.class);

    public transient AggregatingState<Tuple2<Integer, Integer>, Double> aggregatingState;

    @Override
    public void open(Configuration config) {
        AggregatingStateDescriptor<Tuple2<Integer, Integer>, AverageAccumulator, Double> descriptor = new AggregatingStateDescriptor<>("AggregatingState",
                new AggregateFunction<Tuple2<Integer, Integer>, AverageAccumulator, Double>() {

                    @Override
                    public AverageAccumulator createAccumulator() {
                        return new AverageAccumulator();
                    }

                    @Override
                    public AverageAccumulator add(Tuple2<Integer, Integer> value, AverageAccumulator accumulator) {
                        accumulator.add(value.f1);
                        return accumulator;
                    }

                    @Override
                    public Double getResult(AverageAccumulator accumulator) {
                        return accumulator.getLocalValue();
                    }

                    @Override
                    public AverageAccumulator merge(AverageAccumulator a, AverageAccumulator b) {
                        a.merge(b);
                        return a;
                    }
                },
                AverageAccumulator.class);
        aggregatingState = getRuntimeContext().getAggregatingState(descriptor);
    }

    @Override
    public void flatMap(Tuple2<Integer, Integer> input, Collector<Tuple2<Integer, Double>> out) throws Exception {
        aggregatingState.add(input);
        out.collect(Tuple2.of(input.f0, aggregatingState.get()));
    }

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        KeyedStream<Tuple2<Integer, Integer>, Tuple> keyedStream = KeyStateBase.before(env);
        keyedStream.flatMap(new AggregatingStateFlatMap()).print();
        env.execute("Intsmaze AggregatingStateFlatMap");
    }

}
